# 编写训练误差函数，用于计算扩散损失,它需要五个参数：
# model:指以x和t作为输入并产生输出的模型对象或函数
# x_0:表示初始输入（x）状态的张量
# alphas_bar_sqrt:这是一个张量，包含计算中使用的与alpha相关的一些值
# one_minus_alphas_bar_sqrt:这是另一个张量，包含计算中使用的与（1-alpha）相关的值
# n_steps:指定步数的整数
import torch


def diffusion_loss_fn(model, x_0, alphas_bar_sqrt, one_minus_alphas_bar_sqrt, n_steps):
    """对任意时刻t进行采样计算loss"""
    # 根据x_0的形状来确定batch_size
    batch_size = x_0.shape[0]

    """ 对一个batchsize样本生成随机的时刻t，t变得随机分散一些，一个batch size里面覆盖更多的t"""
    # 生成0到n_steps之间的随机整数t（不包括），大小为     （batch_size//2），这些整数表示计算的时间步长
    t = torch.randint(0, n_steps, size=(batch_size // 2,))

    # t的形状（bz） #沿着指定的维度（dim=0）将t与（n_steps-1-t）连接起来，以创建一个以对称顺序表示时间步长的张量
    t = torch.cat([t, n_steps - 1 - t], dim=0)

    # 通过用-1作为参数取消对t的压缩，为它添加了一个额外的维度，使其与后续计算兼容
    t = t.unsqueeze(-1)  # t的形状（bz,1）

    # 使用t对张量alphas_bar_sqrt进行索引，以检索与时间步长相对应的alpha值，即 x0的系数，根号下(alpha_bar_t)
    a = alphas_bar_sqrt[t]

    # 对one_minus_alphas_bar_sqrt执行相同的操作来获得（1-alpha）值，即 eps的系数,根号下(1-alpha_bar_t)
    aml = one_minus_alphas_bar_sqrt[t]

    # 使用torch.randn_like（x_0）生成与x_0形状相同的随机噪音eps
    e = torch.randn_like(x_0)

    # 通过将x_0与a相乘并将e与aml的逐元素相乘来计算x，构造模型的输入
    x = x_0 * a + e * aml

    # 将x和压缩的t（去除额外的维度）送入模型进行计算，得到t时刻的随机噪声预测值
    output = model(x, t.squeeze(-1))

    # 与真实噪声一起计算误差，并返回e和输出之间的均方差，求平均值
    return (e - output).square().mean()


'''diffusion_loss_fn函数生成随机时间步长，使用这些时间步长计算扩散过程，将所得张量传递给指定的模型，并基于随机张量e和模型输出之间的差计算损失'''
